/**
 * Copyright 2024 Bloom Labs Inc
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import type {
  GenerateRequest,
  GenerateResponseData,
  ModelReference,
  StreamingCallback,
} from 'genkit';
import { z } from 'genkit';
import type { GenerateResponseChunkData, ModelAction } from 'genkit/model';
import { modelRef } from 'genkit/model';
import { model } from 'genkit/plugin';

import type { ModelInfo } from 'genkit/model';
import { BetaRunner, Runner } from './runner/index.js';
import {
  AnthropicBaseConfigSchema,
  AnthropicBaseConfigSchemaType,
  AnthropicConfigSchema,
  AnthropicThinkingConfigSchema,
  resolveBetaEnabled,
  type ClaudeModelParams,
  type ClaudeRunnerParams,
} from './types.js';

// This contains all the Anthropic config schema types
type ConfigSchemaType =
  | AnthropicBaseConfigSchemaType
  | AnthropicThinkingConfigSchemaType;

/**
 * Creates a model reference for a Claude model.
 */
function commonRef(
  name: string,
  configSchema: ConfigSchemaType = AnthropicConfigSchema,
  info?: ModelInfo
): ModelReference<ConfigSchemaType> {
  return modelRef({
    name: `anthropic/${name}`,
    configSchema,
    info: info ?? {
      supports: {
        multiturn: true,
        tools: true,
        media: true,
        systemRole: true,
        output: ['text'],
      },
    },
  });
}

export const KNOWN_CLAUDE_MODELS: Record<
  string,
  ModelReference<
    AnthropicBaseConfigSchemaType | AnthropicThinkingConfigSchemaType
  >
> = {
  'claude-3-haiku': commonRef('claude-3-haiku', AnthropicBaseConfigSchema),
  'claude-3-5-haiku': commonRef('claude-3-5-haiku', AnthropicBaseConfigSchema),
  'claude-sonnet-4': commonRef(
    'claude-sonnet-4',
    AnthropicThinkingConfigSchema
  ),
  'claude-opus-4': commonRef('claude-opus-4', AnthropicThinkingConfigSchema),
  'claude-sonnet-4-5': commonRef(
    'claude-sonnet-4-5',
    AnthropicThinkingConfigSchema
  ),
  'claude-haiku-4-5': commonRef(
    'claude-haiku-4-5',
    AnthropicThinkingConfigSchema
  ),
  'claude-opus-4-1': commonRef(
    'claude-opus-4-1',
    AnthropicThinkingConfigSchema
  ),
};

/**
 * Gets the un-prefixed model name from a modelReference.
 */
export function extractVersion(
  model: ModelReference<ConfigSchemaType> | undefined,
  modelName: string
): string {
  // Extract from model name (remove 'anthropic/' prefix if present)
  return modelName.replace(/^anthropic\//, '');
}

/**
 * Generic Claude model info for unknown/unsupported models.
 * Used when a model name is not in KNOWN_CLAUDE_MODELS.
 */
export const GENERIC_CLAUDE_MODEL_INFO = {
  supports: {
    multiturn: true,
    tools: true,
    media: true,
    systemRole: true,
    output: ['text'],
  },
};

export type KnownClaudeModels = keyof typeof KNOWN_CLAUDE_MODELS;
export type ClaudeModelName = string;
export type AnthropicConfigSchemaType = typeof AnthropicConfigSchema;
export type AnthropicThinkingConfigSchemaType =
  typeof AnthropicThinkingConfigSchema;
export type ClaudeConfig = z.infer<typeof AnthropicConfigSchema>;

/**
 * Creates the runner used by Genkit to interact with the Claude model.
 * @param params Configuration for the Claude runner.
 * @param configSchema The config schema for this model (used for type inference).
 * @returns The runner that Genkit will call when the model is invoked.
 */
export function claudeRunner<TConfigSchema extends z.ZodTypeAny>(
  params: ClaudeRunnerParams,
  configSchema: TConfigSchema
) {
  const { defaultApiVersion, ...runnerParams } = params;

  if (!runnerParams.client) {
    throw new Error('Anthropic client is required to create a runner');
  }

  let stableRunner: Runner | null = null;
  let betaRunner: BetaRunner | null = null;

  return async (
    request: GenerateRequest<TConfigSchema>,
    {
      streamingRequested,
      sendChunk,
      abortSignal,
    }: {
      streamingRequested: boolean;
      sendChunk: StreamingCallback<GenerateResponseChunkData>;
      abortSignal: AbortSignal;
    }
  ): Promise<GenerateResponseData> => {
    // Cast to AnthropicConfigSchema for internal runner which expects the full schema
    const normalizedRequest = request as unknown as GenerateRequest<
      typeof AnthropicConfigSchema
    >;
    const isBeta = resolveBetaEnabled(
      normalizedRequest.config,
      defaultApiVersion
    );
    const runner = isBeta
      ? (betaRunner ??= new BetaRunner(runnerParams))
      : (stableRunner ??= new Runner(runnerParams));
    return runner.run(normalizedRequest, {
      streamingRequested,
      sendChunk,
      abortSignal,
    });
  };
}

/**
 * Strips the 'anthropic/' namespace prefix if present.
 */
function checkModelName(name: string): string {
  return name.startsWith('anthropic/') ? name.slice(10) : name;
}

/**
 * Creates a model reference for a Claude model.
 * This allows referencing models without initializing the plugin.
 */
export function claudeModelReference(
  name: string,
  config: z.infer<typeof AnthropicConfigSchema> = {}
): ModelReference<z.ZodTypeAny> {
  const modelName = checkModelName(name);
  return modelRef({
    name: `anthropic/${modelName}`,
    config: config,
    configSchema: AnthropicConfigSchema,
    info: {
      ...GENERIC_CLAUDE_MODEL_INFO,
    },
  });
}

/**
 * Defines a Claude model with the given name and Anthropic client.
 * Accepts any model name and lets the API validate it. If the model is in KNOWN_CLAUDE_MODELS, uses that modelRef
 * for better defaults; otherwise creates a generic model reference.
 */
export function claudeModel(
  params: ClaudeModelParams
): ModelAction<z.ZodTypeAny> {
  const {
    name,
    client: runnerClient,
    cacheSystemPrompt: cachePrompt,
    defaultApiVersion: apiVersion,
  } = params;
  // Use supported model ref if available, otherwise create generic model ref
  const modelRef = KNOWN_CLAUDE_MODELS[name];
  const modelInfo = modelRef ? modelRef.info : GENERIC_CLAUDE_MODEL_INFO;
  const configSchema = modelRef?.configSchema ?? AnthropicConfigSchema;

  return model<
    AnthropicBaseConfigSchemaType | AnthropicThinkingConfigSchemaType
  >(
    {
      name: `anthropic/${name}`,
      ...modelInfo,
      configSchema: configSchema,
    },
    claudeRunner(
      {
        name,
        client: runnerClient,
        cacheSystemPrompt: cachePrompt,
        defaultApiVersion: apiVersion,
      },
      configSchema
    )
  );
}
